# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging
import os

import jinja2
from absl import app, flags
from tornado import gen, web
from tornado.ioloop import IOLoop


def parse_date(date_str):
  try:
    return datetime.datetime.strptime(date_str, '%Y%m%d')
  except ValueError:
    return None


def get_latest_plot_dir(plot_dir):
  dirs = []
  for fn in os.listdir(plot_dir):
    fp = os.path.join(plot_dir, fn)
    if os.path.isdir(fp) and parse_date(fn):
      dirs.append(fn)
  if not dirs:
    raise ValueError('Not exists')
  return max(dirs)


def list_files(target_dir, extension):
  files = []
  if not os.path.isdir(target_dir):
    return files
  for fn in os.listdir(target_dir):
    fp = os.path.join(target_dir, fn)
    if fn.endswith(extension) and os.path.isfile(fp):
      files.append(fn)
  return sorted(files)


class TemplateEnv(object):
  def __init__(self, template_path):
    template_loader = jinja2.FileSystemLoader(searchpath=template_path)
    self._template_env = jinja2.Environment(loader=template_loader)

  def get_template(self, filename):
    return self._template_env.get_template(filename)


class MainHandler(web.RequestHandler):
  def initialize(self, plot_dir, template_env):
    self._plot_dir = plot_dir
    self._template_env = template_env

  @gen.coroutine
  def get(self, path):
    if not path:
      today_str = get_latest_plot_dir(self._plot_dir)
      today = parse_date(today_str)
    else:
      today = parse_date(path)
      if not today:
        self.set_status(400)
        return
      today_str = today.strftime('%Y%m%d')

    files = list_files(os.path.join(self._plot_dir, today_str), '.png')
    template = self._template_env.get_template('index.html.jinja2')

    plots = []
    for fn in files:
      plots.append({'name': fn[:-4], 'url': '/plot/%s/%s' % (today_str, fn)})

    prev_date_str = (today - datetime.timedelta(days=1)).strftime('%Y%m%d')
    next_date_str = (today + datetime.timedelta(days=1)).strftime('%Y%m%d')
    out = template.render(date=today_str,
                          prev_date=prev_date_str,
                          next_date=next_date_str,
                          plots=plots)
    self.write(out)


class NonCacheStaticFileHandler(web.StaticFileHandler):
  def initialize(self, plot_dir, **kwargs):
    self._plot_dir = plot_dir
    super().initialize(**kwargs)

  def set_extra_headers(self, path):
    if path[:8] == get_latest_plot_dir(self._plot_dir):
      # Disable cache
      self.set_header('Cache-Control', 'no-store, no-cache, must-revalidate, max-age=0')


def main(argv):
  FLAGS = flags.FLAGS
  plot_dir = os.path.abspath(os.path.normpath(os.path.expanduser(FLAGS.plot_dir)))
  template_env = TemplateEnv('coin/strategy/tool/monitor/static/')

  app = web.Application([
      (r'/plot/(.*)', NonCacheStaticFileHandler, {
          'plot_dir': plot_dir, 'path': plot_dir
      }), (r'/(.*)', MainHandler, {
          'plot_dir': plot_dir, 'template_env': template_env
      })
  ])
  app.listen(FLAGS.port)
  IOLoop.current().start()


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')

  flags.DEFINE_string('plot_dir', 'plot_out', 'Plot directory')
  flags.DEFINE_integer('port', 8888, 'Port')
  app.run(main)
